package com.tianji.memory.mysql;

import cn.hutool.core.collection.CollStreamUtil;
import cn.hutool.core.convert.Convert;
import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.tianji.entity.ChatRecord;
import com.tianji.memory.util.MessageUtil;
import com.tianji.service.IChatRecordService;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.Message;
import org.springframework.util.Assert;

import java.util.List;

/**
 * 基于 mysql 的 ChatMemoryRepository
 */
public class JdbcChatMemoryRepository implements ChatMemoryRepository {

    @Resource
    private IChatRecordService chatRecordService;

    @Override
    public List<String> findConversationIds() {
        QueryWrapper<ChatRecord> queryWrapper = new QueryWrapper<>();
        queryWrapper.select("distinct conversation_id");
        List<ChatRecord> conversationIds = chatRecordService.list(queryWrapper);
        return CollStreamUtil.toList(conversationIds, ChatRecord::getConversationId);
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        List<ChatRecord> chatRecords = chatRecordService.lambdaQuery()
                .eq(ChatRecord::getConversationId, conversationId)
                //.orderByAsc(ChatRecord::getCreateTime) 
                .list();
        return CollStreamUtil.toList(chatRecords, r -> MessageUtil.toMessage(r.getData()));
    }

    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        Assert.notEmpty(messages, "消息列表不能为空");
        // 先删除
        deleteByConversationId(conversationId);
        // 再保存所有
        Long userId = Convert.toLong(StrUtil.subBefore(conversationId, "_", true));
        List<ChatRecord> toSave = CollStreamUtil.toList(messages, message ->
                ChatRecord.builder()
                        .conversationId(conversationId)
                        .data(MessageUtil.toJson(message)) // 转成我们自定义的 MyMessage 然后 转成json
                        .creater(userId) // 用户id从 conversationId 中截取
                        .updater(userId)
                        .build()
        );
        chatRecordService.saveBatch(toSave);
    }



    @Override
    public void deleteByConversationId(String conversationId) {
        chatRecordService.lambdaUpdate()
                .eq(ChatRecord::getConversationId, conversationId)
                .remove();
    }

}